import math

import torch
import torch.nn as nn

# Input : (batch_size, seq_len)
# Output : (batch_size, seq_len, feature)

class EmbeddingsLayer(nn.Module):
    def __init__(self, vocab_num, d_model):
        super(EmbeddingsLayer, self).__init__()
        # vocab 代表了 可学习的 token 的数量，而 d_model 是输出向量/矩阵的大小
        # embedding 会创建一个 [vocab_num, d_model] 的可学习矩阵
        self.lookup_table = nn.Embedding(vocab_num, d_model)
        self.d_model = d_model

    # 比如 x = [21, 45, 86], loopup_table 会查找 21 对应的 vector，45 对应的 vector 等等返回
    def forward(self, x):
        # 乘以 sqrt(d_model) 是为了平衡后面添加的 position 信息 和 编码信息
        return self.lookup_table(x) * math.sqrt(self.d_model)


if __name__ == '__main__':
    vocab_num  = 100   # 假设词表大小
    d_model    = 512   # 嵌入维度
    batch_size = 2
    seq_len    = 4

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = EmbeddingsLayer(vocab_num, d_model).to(device)
    model.eval()

    # 模拟一个输入 batch，数值必须小于 vocab_num
    x = torch.tensor([[1, 5, 23, 11],
                      [8, 9, 10, 3]]).to(device)

    output = model(x)

    print("Input token IDs:", x)
    print("Input shape:", x.shape)

    print("\nOutput shape:", output.shape)  # 应该是 [batch, seq_len, d_model]
    # 输出第一个 sequence ,就是 [1, 5, 23, 11] 对应的 seq
    print("\nOutput (first sequence):")
    print(output[0])
    print(output[0].shape) # 应该是 [seq_len, d_model]

    # 进一步测试缩放因子 验证输出是不是乘了 sqrt(d_model)
    raw_embedding = model.lookup_table(x[0, 0])  # 拿第一个token的原始embedding
    scaled_embedding = raw_embedding * math.sqrt(d_model)
    print("\nManually scaled first embedding vs model output:")
    # 理论输出
    print("Manual:", scaled_embedding)
    # 实际输出
    print("From model:", output[0, 0])
    # 理论输出和实际输出的 差距
    print("Difference:", torch.abs(scaled_embedding - output[0, 0]).max())


